# 导入必要的库
import os  # 导入 os 模块

import evaluate
import numpy as np
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer

# 定义模型保存目录
output_dir = "../test_trainer"

# 设置环境变量，避免 tokenizers 并行警告
os.environ["TOKENIZERS_PARALLELISM"] = "false"  # 最佳实践：在导入 tokenizer 之前设置

# 加载预训练的分词器
tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-cased")

# 示例：分词和解码
encoded_input = tokenizer("robert boulter is an english film , television and theatre actor .")
print(encoded_input)
print(tokenizer.decode(encoded_input["input_ids"]))

# 加载数据集
dataset = load_dataset("yelp_review_full")
print(dataset["train"][100])


# 定义分词函数
def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)


# 对数据集进行分词
# 修改：使用 num_proc 控制并行进程数, 加速处理。  根据你的 CPU 核心数调整。
tokenized_datasets = dataset.map(tokenize_function, batched=True, num_proc=8)

# 选择小规模数据集用于训练和验证
small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(1000))
small_eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(1000))

# 加载预训练模型（带有分类头）
model = AutoModelForSequenceClassification.from_pretrained("google-bert/bert-base-cased", num_labels=5)

# 训练参数
training_args = TrainingArguments(
    output_dir=output_dir,
    save_strategy="epoch",
    eval_strategy="epoch",
    num_train_epochs=3,
    logging_dir=f"{output_dir}/logs",
    # 添加：per_device_train_batch_size 和 per_device_eval_batch_size
    per_device_train_batch_size=16,  # 根据你的 GPU 内存调整
    per_device_eval_batch_size=64,  # 可以比训练 batch size 大
    # 添加：learning_rate，一个重要的超参数
    learning_rate=5e-5,
    # 添加：weight_decay, 防止过拟合
    weight_decay=0.01,
    # 添加：load_best_model_at_end=True ,自动加载最佳模型
    load_best_model_at_end=True,
    # 添加metric_for_best_model，指定最佳模型的评估指标。
    metric_for_best_model="accuracy",
)

# 评估指标
metric = evaluate.load("accuracy")


def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)


# 初始化 Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=small_train_dataset,
    eval_dataset=small_eval_dataset,
    compute_metrics=compute_metrics,
)

# 开始训练
trainer.train()

# 训练后保存模型和 tokenizer
trainer.save_model(output_dir)  # 保存模型
tokenizer.save_pretrained(output_dir)  # 重新保存 tokenizer，确保不会丢失
